#!/usr/bin/env python3
import yaml
import os
import shutil
from pathlib import Path
import argparse

def rename_paths(yaml_file):
    """重命名现有路径中的 embodiment 名称为 group 名称"""
    # 加载 YAML 文件
    with open(yaml_file, 'r') as f:
        yaml_data = yaml.safe_load(f)
    
    # 创建 embodiment 到 group 的映射
    emb_to_group = {}
    for group_name, group_config in yaml_data['groups'].items():
        embodiment = group_config.get('embodiments')
        if embodiment:
            emb_to_group[embodiment] = group_name
    
    print(f"Embodiment to group mapping: {emb_to_group}")
    
    # 处理所有路径
    for path_type in ['hdf5_dirs', 'zarrs', 'pcds']:
        if path_type not in yaml_data:
            continue
            
        for idx, data in yaml_data[path_type].items():
            old_path = data['name']
            if not os.path.exists(old_path):
                print(f"跳过不存在的路径: {old_path}")
                continue
            print(f"old_path: {old_path}")
            # 检查路径中是否包含 embodiment 名称
            for embodiment, group_name in emb_to_group.items():
                if embodiment in old_path:
                    # 创建新路径 - 直接替换 embodiment 名称为 group 名称
                    new_path = old_path.replace(embodiment, group_name)
                    print(f"old_path: {old_path}, new_path: {new_path}")
                    if new_path != old_path:
                        print(f"重命名: {old_path} -> {new_path}")
                        
                        # 确保目标目录存在
                        new_dir = os.path.dirname(new_path)
                        os.makedirs(new_dir, exist_ok=True)
                        
                        # 重命名文件或目录
                        try:
                            if os.path.isdir(old_path):
                                shutil.move(old_path, new_path)
                            else:
                                shutil.move(old_path, new_path)
                            print(f"✅ 成功重命名: {old_path}")
                            
                            # 更新 YAML 中的路径
                            yaml_data[path_type][idx]['name'] = new_path
                            
                        except Exception as e:
                            print(f"❌ 重命名失败: {old_path} -> {new_path}, 错误: {e}")
                        break


def main():
    parser = argparse.ArgumentParser(description='重命名路径中的 embodiment 名称为 group 名称')
    parser.add_argument('yaml_file', help='YAML 配置文件路径')
    args = parser.parse_args()
    
    if not os.path.exists(args.yaml_file):
        print(f"错误: YAML 文件不存在: {args.yaml_file}")
        return
    
    rename_paths(args.yaml_file)

if __name__ == "__main__":
    main() 